`timescale 1ns/1ps
`include "Func.v"
`define data_width 2 * 12
// input   a*2 b*2 c*2 w*2 
// output  Dout *4 
// NTT:    Dout[0] =  a[0] + b[0]*w[0]  Dout[2] =  a[0] - b[0]*w[0]      Dout[1] =  a[1] + b[1]*w[1]   Dout[3] =  a[1] - b[1]*w[1]
// INTT:   Dout[0] = (a[0] + b[0])/2    Dout[2] = (a[0] - b[0])*w[0]/2   Dout[1] = (a[1] + b[1])/2     Dout[3] = (a[1] - b[1])*w[1]/2
// PWM1:   Dout[0] =  a[0] + b[0]       Dout[2] =  b[0]*w[0]             Dout[1] =  a[1] + b[1]*       Dout[3] =  b[1]*w[1]
// PWM2:   Dout[0] =  a[0] + b[0]*w[0]  Dout[2] =  b[0]*w[0]-a[0]-c[0]   Dout[1] =  a[1] + b[1]*w[1]   Dout[2] =  b[1]*w[1]-a[1]-c[1]
// Encode_ADD : begin o[i] = ma0_s[i]; e[i] = b[i]    ; end
// PolyAdd    : begin o[i] = ma0_s[i]; e[i] = 12'd0   ; end
// Decode_SUB : begin o[i] = 12'd0   ; e[i] = ms0_s[i]; end

module PLU (
    input  wire clk, rst_n,
    input  wire [`data_width - 1 : 0] ain,bin,cin,win,
    input  wire [ 5:0] opcode, 
    output wire [`data_width * 2 - 1: 0] Dout  
);

genvar i ;

wire [11:0] a[0:1];
wire [11:0] b[0:1];
wire [11:0] c[0:1];
wire [11:0] w[0:1];

wire [11:0] encode_add[0:1];        //加密模加 加法器输出
//**************输入转换**********************//
generate
    for(i=0;i<2;i=i+1) begin
        assign a[i]   = ain[12*i+11 : 12*i];
        assign b[i]   = bin[12*i+11 : 12*i];
        assign c[i]   = cin[12*i+11 : 12*i];
        assign w[i]   = win[12*i+11 : 12*i];
    end
endgenerate
     
//add0 输入 输出
wire  [11: 0] ma0_a[0:1] , ma0_s[0:1];
reg   [11: 0] ma0_b[0:1] ;
generate
    for(i=0;i<2;i=i+1) begin
        always@(*) begin
        if(opcode == `Encode_ADD) ma0_b[i] = c[i];  //加密模加
        else if(opcode == `PWM2) begin if(i%2==0) ma0_b[i] = b[i];  else ma0_b[i] = c[i]; end 
        else ma0_b[i] = b[i];                       //其余选择 a[i]
        end

        assign ma0_a[i] = a[i];

        MA_s theMA0(.MA_a(ma0_a[i]),.MA_b(ma0_b[i]),.MA_s(ma0_s[i]));
    end
endgenerate

//sub0 输入 输出
wire [11:0] ms0_a[0:1], ms0_b[0:1];
wire [11:0] ms0_s[0:1];
generate
    for(i=0;i<2;i=i+1) begin
        assign ms0_a[i] = a[i];
        assign ms0_b[i] = b[i];

        MS_s theMS0(.MS_a(ms0_a[i]),.MS_b(ms0_b[i]),.MS_s(ms0_s[i]));
    end    
endgenerate

//第一级 选择器输出  o e 
reg [11:0] o[0:1], e[0:1];
generate
    for(i=0;i<2;i=i+1) begin
        always@(*) begin
            case(opcode) 
            `NTT        : begin o[i] = a[i]    ; e[i] = b[i];     end
            `INTT       : begin o[i] = ma0_s[i]; e[i] = ms0_s[i]; end
            `PWM1       : begin o[i] = ma0_s[i]; e[i] = b[i];     end
            `PWM2       : begin if(i%2==0) begin o[i] =   a[i]  ; e[i] = b[i]; end
                                else       begin o[i] = ma0_s[i]; e[i] = b[i]; end
                          end
            `Encode_ADD : begin o[i] = ma0_s[i]; e[i] = b[i]    ; end
            `PolyAdd    : begin o[i] = ma0_s[i]; e[i] = 12'd0   ; end
            `Decode_SUB : begin o[i] = 12'd0   ; e[i] = ms0_s[i]; end
            default     : begin o[i] = 12'd0   ; e[i] = 12'd0   ; end
            endcase
        end
    end
endgenerate

//o_reg  延迟链
reg [11:0] o_reg[0:1][0:3]; 

generate
    for(i=0;i<2;i=i+1) begin
        always@(posedge clk or negedge rst_n)
        begin
            if(!rst_n) begin o_reg[i][0]<=0; o_reg[i][1]<=0; o_reg[i][2]<=0; o_reg[i][3]<=0;end
            else  begin {o_reg[i][3],o_reg[i][2],o_reg[i][1],o_reg[i][0]}  <= {o_reg[i][2],o_reg[i][1],o_reg[i][0],o[i]}; end 
        end  
        end  
endgenerate

//e_reg 
reg [11:0] e_reg[0:1];
generate
    for(i=0;i<2;i=i+1) begin
        always@(posedge clk or negedge rst_n)
        begin
            if(!rst_n) e_reg[i]<=12'd0;
            else e_reg[i]<=e[i];
        end
    end
endgenerate

//w_reg
reg [11:0] w_reg[0:1];
generate
    for(i=0;i<2;i=i+1)
    begin
        always@(posedge clk or negedge rst_n)
        begin
            if(!rst_n) w_reg[i]<=12'd0;
            else if(opcode == `Encode_ADD) w_reg[i]<=12'd1;
            else w_reg[i]<=w[i];
        end
    end
endgenerate

//MulMod 模乘器
wire [11:0] MUL_a[0:1],MUL_b[0:1];
wire [11:0] P[0:1];
wire Mul_en;
wire mul_mode = opcode[5];
wire Xpress_mode = opcode[3];
wire d_mode = opcode[0];
wire k_mode = opcode[1];
wire [11:0] do_compress[0:1];
wire [11:0] do_decompress[0:1];

assign Mul_en = !(opcode == `PolyAdd || opcode == `Decode_SUB);  
generate
    for(i=0;i<2;i=i+1)
    begin
        assign MUL_a[i] = e_reg[i];
        assign MUL_b[i] = w_reg[i];
        MUL_MOD theMulMod(.clk(clk),.en(Mul_en),.rst_n(rst_n),.a(MUL_a[i]),.b(MUL_b[i]),.product(P[i]),
                            .mul_mode(mul_mode),.Xpress_mode(Xpress_mode),.d_mode(d_mode),.k_mode(k_mode),.din_compress(a[i]),.do_compress(do_compress[i]),.din_decompress(a[i]),.do_decompress(do_decompress[i])); /////////////////////////////////////
    end
endgenerate

//  mul_mode, // mul mode select, 0-as MUL MOD, 1 as-compress/decompress
//  Xpress_mode, // 0-compress, 1-decompress
//  d_mode, // d_mode, 0 = du, 1 = dv
//  k_mode, // k_mode, 0 = kyber512/768, 1 = kyber1024
//                     // for d_mode and k_mode select, the state is below 
//                     //        0         0            du10,  du=10 kyber512/768
//                     //        1         0            dv4,   dv=4  kyber512/768
//                     //        0         1            du11,  du=11 kyber1024
//                     //        1         1            dv5,   dv=5  kyber1024

//add1 输入输出
wire [11: 0] ma1_a[0:1] ;
wire [11: 0] ma1_b[0:1] , ma1_s[0:1];
generate
    for(i=0;i<2;i=i+1) begin
        assign ma1_a[i] = o_reg[i][3];
        assign ma1_b[i] = P[i];

        MA_s theMA1(.MA_a(ma1_a[i]),.MA_b(ma1_b[i]),.MA_s(ma1_s[i]));
    end
endgenerate

//*****************  sub1 输入 输出 ************************//
wire [11:0] ms1_a[0:1], ms1_b[0:1];
wire [11:0] ms1_s[0:1];
generate
    for(i=0;i<2;i=i+1) begin
        assign ms1_a[i] = o_reg[i][3];
        assign ms1_b[i] = P[i];

        MS_s theMS1(.MS_a(ms1_a[i]),.MS_b(ms1_b[i]),.MS_s(ms1_s[i]));
    end    
endgenerate

//******************** div2 ************************//
wire [11:0] P_div2[0:1];
wire [11:0] sum_div2[0:1];

generate
    for(i=0;i<2;i=i+1)
    begin
        Div2 the_p_div2  (.Din(P[i])        ,.Dout(P_div2[i]) );
        Div2 the_sum_div2(.Din(o_reg[i][3]) ,.Dout(sum_div2[i]));
    end
endgenerate

//******************** out *************************//

reg [11:0] out[0:3];
generate
    for(i=0;i<2;i=i+1) begin
        always@(*)
        begin
            case(opcode)
            `NTT        : begin out[i] = ma1_s[i]   ; out[i+2] = ms1_s[i]   ; end
            `INTT       : begin out[i] = sum_div2[i]; out[i+2] = P_div2[i]  ; end
            `PWM1       : begin out[i] = o_reg[i][3]; out[i+2] = P[i]       ; end
            `PWM2       : begin out[i] = ma1_s[i]   ; out[i+2] = 12'd3329 - ms1_s[i]; end
            `Encode_ADD : begin out[i] = ma1_s[i]   ; out[i+2] = 12'd0;   end
            `PolyAdd    : begin out[i] = o_reg[i][0]; out[i+2] = 12'd0;   end
            `Decode_SUB : begin out[i] = e_reg[i]   ; out[i+2] = 12'd0;   end
            
            `Compress10,`Compress4,`Compress11,`Compress5:  begin out[i] = do_compress[i] ; out[i+2] =12'd0; end
//            `Compare10,`Compare4,`Compare11,`Compare5:      begin out[i] = (do_compress[i] == b[i]) ; out[i+2] =12'd0; end
            `Decompress10,`Decompress4,`Decompress11,`Decompress5: begin out[i] = do_decompress[i] ; out[i+2] =12'd0; end
            default     : begin out[i] = 12'd0      ; out[i+2] = 12'd0;   end
            endcase
        end
    end
endgenerate

//****************输出转换****************************//
generate
    for(i=0;i<2;i=i+1) begin
        assign Dout[12*i+11 : 12*i   ] = out[i]   ;
        assign Dout[12*i+35 : 12*i+24] = out[i+2] ;
    end
endgenerate


//***************************************************//


endmodule


module Alu_tb;

reg clk, rst_n;
reg   [23:0] ain;
reg   [23:0] bin;
reg   [23:0] cin;
reg   [23:0] win;
wire  [47:0] out;
reg [5:0]  opcode;


wire  [23:0] ct_check[0:3];
wire  [23:0] PWM1_check[0:3];
wire  [23:0] PWM2_check[0:3];
wire  [23:0] P_check[0:1];
wire  [23:0] gs[0:3];
wire  [23:0] gs_check[0:3];
ALU theAlu(
    .clk(clk),.rst_n(rst_n),
    .ain(ain),.bin(bin),.win(win), .cin(cin),
    .opcode(opcode),  //000：NTT正变换  001：NTT逆变换  010：点乘 8项乘法  011：点乘 四项加法  100：a+bw  ----2  a+b -----2 
    .Dout(out));


wire [11:0] a[0:1];
wire [11:0] b[0:1];
wire [11:0] c[0:1];
wire [11:0] w[0:1];

wire [11:0] encode_add[0:1];        //加密模加 加法器输出

genvar i;
//**************输入转换**********************//
generate
    for(i=0;i<2;i=i+1) begin
        assign a[i]   = ain[12*i+11 : 12*i];
        assign b[i]   = bin[12*i+11 : 12*i];
        assign c[i]   = cin[12*i+11 : 12*i];
        assign w[i]   = win[12*i+11 : 12*i];
    end
endgenerate



initial begin
    rst_n=0;
    clk=0;
    opcode =0;
    #20 rst_n =0;
    #10 rst_n =1;
    #1000 opcode =`NTT;
    #1000 opcode =`INTT;
    #1000 opcode =`PWM1;
    #1000 opcode =`PWM2;
    #1000 opcode = `Compress4;
end

always #5 clk=~clk;


generate
    for(i=0;i<2;i=i+1) begin
        always@(posedge clk) begin
            ain[12*i + 11 : 12 * i]<={$random}%3329;
            bin[12*i + 11 : 12 * i]<={$random}%3329;
            cin[12*i + 11 : 12 * i]<={$random}%3329;
            win[12*i + 11 : 12 * i]<={$random}%3329;
        end
    end
endgenerate

wire [11:0] Dout[0:3];
generate
    for(i=0;i<2;i=i+1) begin
        assign Dout[i]   = out[12*i+11 : 12*i   ];
        assign Dout[i+2] = out[12*i+35 : 12*i+24];
    end
endgenerate

generate
    for(i=0;i<2;i=i+1)
    begin
        assign ct_check[i]     = (a[i] + b[i] * w[i] % 3329)%3329;
        assign ct_check[i+2]   = (a[i] - b[i] * w[i] % 3329 + 3329 )%3329; 
        assign gs[i]           = (a[i] + b[i])%3329;
        assign gs[i+2]         = ((a[i] + 3329 - b[i])* w[i])%3329;
        assign gs_check[i]     = gs[i][0]  ? (  gs[i][11:1]+12'd1665) :   gs[i][11:1];
        assign gs_check[i+2]   = gs[i+2][0]? (gs[i+2][11:1]+12'd1665) : gs[i+2][11:1];
        assign PWM1_check[i]   = (a[i] + b[i]) %3329 ; 
        assign PWM1_check[i+2] = (b[i] * w[i]) %3329 ;
        assign PWM2_check[i]   = (a[i] + b[i] * w[i] % 3329)%3329;
        assign PWM2_check[i+2] = (b[i] * w[i] % 3329 + 2*3329 - a[i] -c[i])%3329;
        assign P_check[i] = (b[i] * w[i]) %3329 ;
    end
endgenerate

endmodule